#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import asyncio
import json
import sys
from typing import List, Optional

from src.agents.AbstractGeneratorAgent import AbstractGeneratorAgent
from src.agents.KeywordExtractorAgent import KeywordExtractorAgent
from src.agents.LiteratureRetrieverAgent import LiteratureRetrieverAgent
from src.agents.OutlineBuilderAgent import OutlineBuilderAgent
from src.agents.CaseAnalyzerAgent import CaseAnalyzerAgent
from src.agents.PolicyEvaluatorAgent import PolicyEvaluatorAgent
from src.agents.base_config import StateManager
from src.utils.env_config import get_task_name


class TaskOrchestrator2:
    """增强型主控智能体：支持动态任务调整与上下文保持"""

    def __init__(self):
        self.state_manager = StateManager()
        self.keyword_agent = KeywordExtractorAgent()
        self.literature_agent = LiteratureRetrieverAgent()
        self.abstract_agent = AbstractGeneratorAgent()
        self.outline_agent = OutlineBuilderAgent()
        self.case_agent = CaseAnalyzerAgent()
        self.policy_agent = PolicyEvaluatorAgent()

    @staticmethod
    def should_case_analysis(retrieval_summary: str) -> bool:
        if retrieval_summary.find("技术突破") >= 0 and retrieval_summary.find("创新案例") >= 0:
            return True
        return False
    @staticmethod
    def should_policy_evaluation(retrieval_summary: str) -> bool:
        if retrieval_summary.find("政策法规") >= 0:
            return True
        return False

    async def execute_task(self, topic: str) -> None:
        """执行完整的报告生成任务（支持动态流程调整）"""
        print(f"开始执行增强型报告生成任务：{topic}")

        # 保存主题到上下文
        self.state_manager.save_context("topic", topic)
        # T1: 关键词提取
        await self._execute_subtask("T1", self._task1_keyword_extraction, topic)

        # T2: 文献检索
        keywords = self.state_manager.get_context("keywords")
        await self._execute_subtask("T2", self._task2_literature_retrieval, keywords)

        # 决策分析：确定后续任务流程
        retrieval_summary = self.state_manager.get_context("retrieval_summary")
        if self.should_case_analysis(retrieval_summary):
            # T5: 技术案例分析
            await self._execute_subtask("T5", self._task5_case_analysis)

        if self.should_policy_evaluation(retrieval_summary):
            # T6: 政策影响评估
            await self._execute_subtask("T6", self._task6_policy_evaluation)


        # T3: 摘要生成
        await self._execute_subtask("T3", self._task3_abstract_generation, retrieval_summary)
        # T4: 报告框架搭建
        abstract = self.state_manager.get_context("abstract")
        await self._execute_subtask("T4", self._task4_outline_building, topic,abstract )

        print("=== 所有任务执行完成！ ===")

    async def _execute_subtask(self, task_id: str, task_func, *args) -> None:
        """执行子任务"""
        task_name = get_task_name(task_id)
        self.state_manager.update_task_status(task_id, "未开始")
        print(f"\n开始执行 {task_id}: {task_name}")
        # 更新状态为"进行中"
        self.state_manager.update_task_status(task_id, "进行中")
        try:
            # 执行任务
            await task_func(*args)
            # 更新状态为"完成"
            self.state_manager.update_task_status(task_id, "完成")
            print(f"{task_id}: {task_name} 执行完成")

        except Exception as e:
            print(f"{task_id}: {task_name} 执行失败: {e}")
            self.state_manager.update_task_status(task_id, "失败")
            raise


    async def _task1_keyword_extraction(self, topic: str) -> None:
        """T1: 关键词提取"""
        keywords = await self.keyword_agent.extract_keywords(topic)
        self.state_manager.save_context("keywords", keywords)
        print(f"提取的关键词: {keywords}")

    async def _task2_literature_retrieval(self, keywords: List[str]) -> None:
        """T2: 文献检索"""
        retrieval_summary = await self.literature_agent.retrieve_literature(keywords)
        self.state_manager.save_context("retrieval_summary", retrieval_summary)
        print(f"文献检索摘要: {retrieval_summary[:100]}...")

    async def _task3_abstract_generation(self, retrieval_summary: str) -> None:
        """T3: 摘要生成"""
        policy_impact = self.state_manager.get_context("policy_impact")
        case_analysis = self.state_manager.get_context("case_analysis")
        if policy_impact:
            retrieval_summary += f"\n\n## 政策影响评估\n{policy_impact}"
        if case_analysis:
            retrieval_summary += f"\n\n## 技术案例分析\n{case_analysis}"

        abstract = await self.abstract_agent.generate_abstract(retrieval_summary)
        self.state_manager.save_context("abstract", abstract)
        print(f"生成的摘要: {abstract[:100]}...")

    async def _task4_outline_building(self, topic: str, abstract: str) -> None:
        """T4: 报告框架搭建"""
        markdown_content = await self.outline_agent.build_outline(topic, abstract)

        # 保存到文件
        with open("report_framework.md", "w", encoding="utf-8") as f:
            f.write(markdown_content)

        self.state_manager.save_context("markdown_content", markdown_content)
        print("报告框架已生成并保存到 report_framework.md")

    async def _task5_case_analysis(self) -> None:
        """T5: 技术案例分析"""
        retrieval_summary = self.state_manager.get_context("retrieval_summary")
        keywords = self.state_manager.get_context("keywords")
        abstract = self.state_manager.get_context("abstract")

        case_analysis = await self.case_agent.analyze_technical_cases(
            retrieval_summary, keywords, abstract
        )
        self.state_manager.save_context("case_analysis", case_analysis)
        print(f"技术案例分析: {case_analysis[:100]}...")

    async def _task6_policy_evaluation(self) -> None:
        """T6: 政策影响评估"""
        retrieval_summary = self.state_manager.get_context("retrieval_summary")
        keywords = self.state_manager.get_context("keywords")
        abstract = self.state_manager.get_context("abstract")

        policy_impact = await self.policy_agent.evaluate_policy_impact(
            retrieval_summary, keywords, abstract
        )
        self.state_manager.save_context("policy_impact", policy_impact)
        print(f"政策影响评估: {policy_impact[:100]}...")


    async def _handle_insufficient_literature(self, topic: str) -> None:
        """处理文献不足场景"""
        print("=== 处理文献不足场景 ===")
        
        # 标记跳过的任务
        for task_id in ["T3", "T5", "T6"]:
            self.state_manager.update_task_status(task_id, "跳过")
        
        # 生成提示性报告框架
        insufficient_content = f"""
        # {topic} 报告
        ## 摘要
        文献不足，建议补充检索。
        ## 分点提纲
        （无）"""
        
        # 保存到文件
        with open("report_framework.md", "w", encoding="utf-8") as f:
            f.write(insufficient_content)
        
        self.state_manager.save_context("markdown_content", insufficient_content)
        self.state_manager.update_task_status("T4", "完成")
        print("已生成提示性报告框架")


async def main():
    # """主函数"""
    # if len(sys.argv) != 2:
    #     print("使用方法: python task1_main.py '报告主题'")
    #     print("示例: python task1_agentscope.py '2025 Q3 AI 行业趋势'")
    #     return  '2025 Q3 AI大模型安全技术突破轻量化部署创新案例政策法规'
    # topic = sys.argv[1]
    topic = '2025 Q3 AI大模型安全政策法规'
    # 创建任务编排器并执行任务
    orchestrator = TaskOrchestrator2()
    await orchestrator.execute_task(topic)


if __name__ == "__main__":
    asyncio.run(main())